{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BERT for Multi-Label Classification\n",
    "Refer post : https://medium.com/@javaid.nabi/building-a-multi-label-text-classifier-using-bert-and-tensorflow-f188e0ecdc5d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: Logging before flag parsing goes to stderr.\n",
      "W0511 13:58:07.529313  4008 __init__.py:56] Some hub symbols are not available because TensorFlow version is less than 1.14\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import collections\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "import tensorflow_hub as hub\n",
    "from datetime import datetime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "##install bert if not already done\n",
    "##!pip install bert-tensorflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import bert\n",
    "from bert import run_classifier\n",
    "from bert import optimization\n",
    "from bert import tokenization\n",
    "from bert import modeling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "##use downloaded model, change path accordingly\n",
    "BERT_VOCAB= './uncased_L-12_H-768_A-12/vocab.txt'\n",
    "BERT_INIT_CHKPNT = './uncased_L-12_H-768_A-12/bert_model.ckpt'\n",
    "BERT_CONFIG = './uncased_L-12_H-768_A-12/bert_config.json'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenization.validate_case_matches_checkpoint(True,BERT_INIT_CHKPNT)\n",
    "tokenizer = tokenization.FullTokenizer(\n",
    "      vocab_file=BERT_VOCAB, do_lower_case=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['this',\n",
       " 'here',\n",
       " \"'\",\n",
       " 's',\n",
       " 'an',\n",
       " 'example',\n",
       " 'of',\n",
       " 'using',\n",
       " 'the',\n",
       " 'bert',\n",
       " 'token',\n",
       " '##izer']"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "tokenizer.tokenize(\"This here's an example of using the BERT tokenizer\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "##change path accordingly\n",
    "train_data_path='./Downloads/train.csv'\n",
    "train = pd.read_csv(train_data_path)\n",
    "test = pd.read_csv('./Downloads/test.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>comment_text</th>\n",
       "      <th>toxic</th>\n",
       "      <th>severe_toxic</th>\n",
       "      <th>obscene</th>\n",
       "      <th>threat</th>\n",
       "      <th>insult</th>\n",
       "      <th>identity_hate</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0000997932d777bf</td>\n",
       "      <td>Explanation\\nWhy the edits made under my usern...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>000103f0d9cfb60f</td>\n",
       "      <td>D'aww! He matches this background colour I'm s...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>000113f07ec002fd</td>\n",
       "      <td>Hey man, I'm really not trying to edit war. It...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0001b41b1c6bb37e</td>\n",
       "      <td>\"\\nMore\\nI can't make any real suggestions on ...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0001d958c54c6e35</td>\n",
       "      <td>You, sir, are my hero. Any chance you remember...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                 id                                       comment_text  toxic  \\\n",
       "0  0000997932d777bf  Explanation\\nWhy the edits made under my usern...      0   \n",
       "1  000103f0d9cfb60f  D'aww! He matches this background colour I'm s...      0   \n",
       "2  000113f07ec002fd  Hey man, I'm really not trying to edit war. It...      0   \n",
       "3  0001b41b1c6bb37e  \"\\nMore\\nI can't make any real suggestions on ...      0   \n",
       "4  0001d958c54c6e35  You, sir, are my hero. Any chance you remember...      0   \n",
       "\n",
       "   severe_toxic  obscene  threat  insult  identity_hate  \n",
       "0             0        0       0       0              0  \n",
       "1             0        0       0       0              0  \n",
       "2             0        0       0       0              0  \n",
       "3             0        0       0       0              0  \n",
       "4             0        0       0       0              0  "
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "ID = 'id'\n",
    "DATA_COLUMN = 'comment_text'\n",
    "LABEL_COLUMNS = ['toxic','severe_toxic','obscene','threat','insult','identity_hate']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class InputExample(object):\n",
    "    \"\"\"A single training/test example for simple sequence classification.\"\"\"\n",
    "\n",
    "    def __init__(self, guid, text_a, text_b=None, labels=None):\n",
    "        \"\"\"Constructs a InputExample.\n",
    "\n",
    "        Args:\n",
    "            guid: Unique id for the example.\n",
    "            text_a: string. The untokenized text of the first sequence. For single\n",
    "            sequence tasks, only this sequence must be specified.\n",
    "            text_b: (Optional) string. The untokenized text of the second sequence.\n",
    "            Only must be specified for sequence pair tasks.\n",
    "            labels: (Optional) [string]. The label of the example. This should be\n",
    "            specified for train and dev examples, but not for test examples.\n",
    "        \"\"\"\n",
    "        self.guid = guid\n",
    "        self.text_a = text_a\n",
    "        self.text_b = text_b\n",
    "        self.labels = labels\n",
    "\n",
    "\n",
    "class InputFeatures(object):\n",
    "    \"\"\"A single set of features of data.\"\"\"\n",
    "\n",
    "    def __init__(self, input_ids, input_mask, segment_ids, label_ids, is_real_example=True):\n",
    "        self.input_ids = input_ids\n",
    "        self.input_mask = input_mask\n",
    "        self.segment_ids = segment_ids\n",
    "        self.label_ids = label_ids,\n",
    "        self.is_real_example=is_real_example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_examples(df, labels_available=True):\n",
    "    \"\"\"Creates examples for the training and dev sets.\"\"\"\n",
    "    examples = []\n",
    "    for (i, row) in enumerate(df.values):\n",
    "        guid = row[0]\n",
    "        text_a = row[1]\n",
    "        if labels_available:\n",
    "            labels = row[2:]\n",
    "        else:\n",
    "            labels = [0,0,0,0,0,0]\n",
    "        examples.append(\n",
    "            InputExample(guid=guid, text_a=text_a, labels=labels))\n",
    "    return examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAIN_VAL_RATIO = 0.9\n",
    "LEN = train.shape[0]\n",
    "SIZE_TRAIN = int(TRAIN_VAL_RATIO*LEN)\n",
    "\n",
    "x_train = train[:SIZE_TRAIN]\n",
    "x_val = train[SIZE_TRAIN:]\n",
    "\n",
    "train_examples = create_examples(x_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_examples_to_features(examples,  max_seq_length, tokenizer):\n",
    "    \"\"\"Loads a data file into a list of `InputBatch`s.\"\"\"\n",
    "\n",
    "    features = []\n",
    "    for (ex_index, example) in enumerate(examples):\n",
    "        print(example.text_a)\n",
    "        tokens_a = tokenizer.tokenize(example.text_a)\n",
    "\n",
    "        tokens_b = None\n",
    "        if example.text_b:\n",
    "            tokens_b = tokenizer.tokenize(example.text_b)\n",
    "            # Modifies `tokens_a` and `tokens_b` in place so that the total\n",
    "            # length is less than the specified length.\n",
    "            # Account for [CLS], [SEP], [SEP] with \"- 3\"\n",
    "            _truncate_seq_pair(tokens_a, tokens_b, max_seq_length - 3)\n",
    "        else:\n",
    "            # Account for [CLS] and [SEP] with \"- 2\"\n",
    "            if len(tokens_a) > max_seq_length - 2:\n",
    "                tokens_a = tokens_a[:(max_seq_length - 2)]\n",
    "\n",
    "        # The convention in BERT is:\n",
    "        # (a) For sequence pairs:\n",
    "        #  tokens:   [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]\n",
    "        #  type_ids: 0   0  0    0    0     0       0 0    1  1  1  1   1 1\n",
    "        # (b) For single sequences:\n",
    "        #  tokens:   [CLS] the dog is hairy . [SEP]\n",
    "        #  type_ids: 0   0   0   0  0     0 0\n",
    "        #\n",
    "        # Where \"type_ids\" are used to indicate whether this is the first\n",
    "        # sequence or the second sequence. The embedding vectors for `type=0` and\n",
    "        # `type=1` were learned during pre-training and are added to the wordpiece\n",
    "        # embedding vector (and position vector). This is not *strictly* necessary\n",
    "        # since the [SEP] token unambigiously separates the sequences, but it makes\n",
    "        # it easier for the model to learn the concept of sequences.\n",
    "        #\n",
    "        # For classification tasks, the first vector (corresponding to [CLS]) is\n",
    "        # used as as the \"sentence vector\". Note that this only makes sense because\n",
    "        # the entire model is fine-tuned.\n",
    "        tokens = [\"[CLS]\"] + tokens_a + [\"[SEP]\"]\n",
    "        segment_ids = [0] * len(tokens)\n",
    "\n",
    "        if tokens_b:\n",
    "            tokens += tokens_b + [\"[SEP]\"]\n",
    "            segment_ids += [1] * (len(tokens_b) + 1)\n",
    "\n",
    "        input_ids = tokenizer.convert_tokens_to_ids(tokens)\n",
    "\n",
    "        # The mask has 1 for real tokens and 0 for padding tokens. Only real\n",
    "        # tokens are attended to.\n",
    "        input_mask = [1] * len(input_ids)\n",
    "\n",
    "        # Zero-pad up to the sequence length.\n",
    "        padding = [0] * (max_seq_length - len(input_ids))\n",
    "        input_ids += padding\n",
    "        input_mask += padding\n",
    "        segment_ids += padding\n",
    "\n",
    "        assert len(input_ids) == max_seq_length\n",
    "        assert len(input_mask) == max_seq_length\n",
    "        assert len(segment_ids) == max_seq_length\n",
    "        \n",
    "        labels_ids = []\n",
    "        for label in example.labels:\n",
    "            labels_ids.append(int(label))\n",
    "\n",
    "        if ex_index < 0:\n",
    "            logger.info(\"*** Example ***\")\n",
    "            logger.info(\"guid: %s\" % (example.guid))\n",
    "            logger.info(\"tokens: %s\" % \" \".join(\n",
    "                    [str(x) for x in tokens]))\n",
    "            logger.info(\"input_ids: %s\" % \" \".join([str(x) for x in input_ids]))\n",
    "            logger.info(\"input_mask: %s\" % \" \".join([str(x) for x in input_mask]))\n",
    "            logger.info(\n",
    "                    \"segment_ids: %s\" % \" \".join([str(x) for x in segment_ids]))\n",
    "            logger.info(\"label: %s (id = %s)\" % (example.labels, labels_ids))\n",
    "\n",
    "        features.append(\n",
    "                InputFeatures(input_ids=input_ids,\n",
    "                              input_mask=input_mask,\n",
    "                              segment_ids=segment_ids,\n",
    "                              label_ids=labels_ids))\n",
    "    return features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We'll set sequences to be at most 128 tokens long.\n",
    "MAX_SEQ_LENGTH = 128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute train and warmup steps from batch size\n",
    "# These hyperparameters are copied from this colab notebook (https://colab.sandbox.google.com/github/tensorflow/tpu/blob/master/tools/colab/bert_finetuning_with_cloud_tpus.ipynb)\n",
    "BATCH_SIZE = 32\n",
    "LEARNING_RATE = 2e-5\n",
    "NUM_TRAIN_EPOCHS = 1.0\n",
    "# Warmup is a period of time where hte learning rate \n",
    "# is small and gradually increases--usually helps training.\n",
    "WARMUP_PROPORTION = 0.1\n",
    "# Model configs\n",
    "SAVE_CHECKPOINTS_STEPS = 1000\n",
    "SAVE_SUMMARY_STEPS = 500"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PaddingInputExample(object):\n",
    "    \"\"\"Fake example so the num input examples is a multiple of the batch size.\n",
    "    When running eval/predict on the TPU, we need to pad the number of examples\n",
    "    to be a multiple of the batch size, because the TPU requires a fixed batch\n",
    "    size. The alternative is to drop the last batch, which is bad because it means\n",
    "    the entire output data won't be generated.\n",
    "    We use this class instead of `None` because treating `None` as padding\n",
    "    battches could cause silent errors.\n",
    "    \"\"\"\n",
    "    \n",
    "    \n",
    "def convert_single_example(ex_index, example, max_seq_length,\n",
    "                           tokenizer):\n",
    "    \"\"\"Converts a single `InputExample` into a single `InputFeatures`.\"\"\"\n",
    "\n",
    "    if isinstance(example, PaddingInputExample):\n",
    "        return InputFeatures(\n",
    "            input_ids=[0] * max_seq_length,\n",
    "            input_mask=[0] * max_seq_length,\n",
    "            segment_ids=[0] * max_seq_length,\n",
    "            label_ids=0,\n",
    "            is_real_example=False)\n",
    "\n",
    "    tokens_a = tokenizer.tokenize(example.text_a)\n",
    "    tokens_b = None\n",
    "    if example.text_b:\n",
    "        tokens_b = tokenizer.tokenize(example.text_b)\n",
    "\n",
    "    if tokens_b:\n",
    "        # Modifies `tokens_a` and `tokens_b` in place so that the total\n",
    "        # length is less than the specified length.\n",
    "        # Account for [CLS], [SEP], [SEP] with \"- 3\"\n",
    "        _truncate_seq_pair(tokens_a, tokens_b, max_seq_length - 3)\n",
    "    else:\n",
    "        # Account for [CLS] and [SEP] with \"- 2\"\n",
    "        if len(tokens_a) > max_seq_length - 2:\n",
    "            tokens_a = tokens_a[0:(max_seq_length - 2)]\n",
    "\n",
    "    # The convention in BERT is:\n",
    "    # (a) For sequence pairs:\n",
    "    #  tokens:   [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]\n",
    "    #  type_ids: 0     0  0    0    0     0       0 0     1  1  1  1   1 1\n",
    "    # (b) For single sequences:\n",
    "    #  tokens:   [CLS] the dog is hairy . [SEP]\n",
    "    #  type_ids: 0     0   0   0  0     0 0\n",
    "    #\n",
    "    # Where \"type_ids\" are used to indicate whether this is the first\n",
    "    # sequence or the second sequence. The embedding vectors for `type=0` and\n",
    "    # `type=1` were learned during pre-training and are added to the wordpiece\n",
    "    # embedding vector (and position vector). This is not *strictly* necessary\n",
    "    # since the [SEP] token unambiguously separates the sequences, but it makes\n",
    "    # it easier for the model to learn the concept of sequences.\n",
    "    #\n",
    "    # For classification tasks, the first vector (corresponding to [CLS]) is\n",
    "    # used as the \"sentence vector\". Note that this only makes sense because\n",
    "    # the entire model is fine-tuned.\n",
    "    tokens = []\n",
    "    segment_ids = []\n",
    "    tokens.append(\"[CLS]\")\n",
    "    segment_ids.append(0)\n",
    "    for token in tokens_a:\n",
    "        tokens.append(token)\n",
    "        segment_ids.append(0)\n",
    "    tokens.append(\"[SEP]\")\n",
    "    segment_ids.append(0)\n",
    "\n",
    "    if tokens_b:\n",
    "        for token in tokens_b:\n",
    "            tokens.append(token)\n",
    "            segment_ids.append(1)\n",
    "        tokens.append(\"[SEP]\")\n",
    "        segment_ids.append(1)\n",
    "\n",
    "    input_ids = tokenizer.convert_tokens_to_ids(tokens)\n",
    "\n",
    "    # The mask has 1 for real tokens and 0 for padding tokens. Only real\n",
    "    # tokens are attended to.\n",
    "    input_mask = [1] * len(input_ids)\n",
    "\n",
    "    # Zero-pad up to the sequence length.\n",
    "    while len(input_ids) < max_seq_length:\n",
    "        input_ids.append(0)\n",
    "        input_mask.append(0)\n",
    "        segment_ids.append(0)\n",
    "\n",
    "    assert len(input_ids) == max_seq_length\n",
    "    assert len(input_mask) == max_seq_length\n",
    "    assert len(segment_ids) == max_seq_length\n",
    "\n",
    "    labels_ids = []\n",
    "    for label in example.labels:\n",
    "        labels_ids.append(int(label))\n",
    "\n",
    "\n",
    "    feature = InputFeatures(\n",
    "        input_ids=input_ids,\n",
    "        input_mask=input_mask,\n",
    "        segment_ids=segment_ids,\n",
    "        label_ids=labels_ids,\n",
    "        is_real_example=True)\n",
    "    return feature\n",
    "\n",
    "\n",
    "def file_based_convert_examples_to_features(\n",
    "        examples, max_seq_length, tokenizer, output_file):\n",
    "    \"\"\"Convert a set of `InputExample`s to a TFRecord file.\"\"\"\n",
    "\n",
    "    writer = tf.python_io.TFRecordWriter(output_file)\n",
    "\n",
    "    for (ex_index, example) in enumerate(examples):\n",
    "        #if ex_index % 10000 == 0:\n",
    "            #tf.logging.info(\"Writing example %d of %d\" % (ex_index, len(examples)))\n",
    "\n",
    "        feature = convert_single_example(ex_index, example,\n",
    "                                         max_seq_length, tokenizer)\n",
    "\n",
    "        def create_int_feature(values):\n",
    "            f = tf.train.Feature(int64_list=tf.train.Int64List(value=list(values)))\n",
    "            return f\n",
    "\n",
    "        features = collections.OrderedDict()\n",
    "        features[\"input_ids\"] = create_int_feature(feature.input_ids)\n",
    "        features[\"input_mask\"] = create_int_feature(feature.input_mask)\n",
    "        features[\"segment_ids\"] = create_int_feature(feature.segment_ids)\n",
    "        features[\"is_real_example\"] = create_int_feature(\n",
    "            [int(feature.is_real_example)])\n",
    "        if isinstance(feature.label_ids, list):\n",
    "            label_ids = feature.label_ids\n",
    "        else:\n",
    "            label_ids = feature.label_ids[0]\n",
    "        features[\"label_ids\"] = create_int_feature(label_ids)\n",
    "\n",
    "        tf_example = tf.train.Example(features=tf.train.Features(feature=features))\n",
    "        writer.write(tf_example.SerializeToString())\n",
    "    writer.close()\n",
    "\n",
    "\n",
    "def file_based_input_fn_builder(input_file, seq_length, is_training,\n",
    "                                drop_remainder):\n",
    "    \"\"\"Creates an `input_fn` closure to be passed to TPUEstimator.\"\"\"\n",
    "\n",
    "    name_to_features = {\n",
    "        \"input_ids\": tf.FixedLenFeature([seq_length], tf.int64),\n",
    "        \"input_mask\": tf.FixedLenFeature([seq_length], tf.int64),\n",
    "        \"segment_ids\": tf.FixedLenFeature([seq_length], tf.int64),\n",
    "        \"label_ids\": tf.FixedLenFeature([6], tf.int64),\n",
    "        \"is_real_example\": tf.FixedLenFeature([], tf.int64),\n",
    "    }\n",
    "\n",
    "    def _decode_record(record, name_to_features):\n",
    "        \"\"\"Decodes a record to a TensorFlow example.\"\"\"\n",
    "        example = tf.parse_single_example(record, name_to_features)\n",
    "\n",
    "        # tf.Example only supports tf.int64, but the TPU only supports tf.int32.\n",
    "        # So cast all int64 to int32.\n",
    "        for name in list(example.keys()):\n",
    "            t = example[name]\n",
    "            if t.dtype == tf.int64:\n",
    "                t = tf.to_int32(t)\n",
    "            example[name] = t\n",
    "\n",
    "        return example\n",
    "\n",
    "    def input_fn(params):\n",
    "        \"\"\"The actual input function.\"\"\"\n",
    "        batch_size = params[\"batch_size\"]\n",
    "\n",
    "        # For training, we want a lot of parallel reading and shuffling.\n",
    "        # For eval, we want no shuffling and parallel reading doesn't matter.\n",
    "        d = tf.data.TFRecordDataset(input_file)\n",
    "        if is_training:\n",
    "            d = d.repeat()\n",
    "            d = d.shuffle(buffer_size=100)\n",
    "\n",
    "        d = d.apply(\n",
    "            tf.contrib.data.map_and_batch(\n",
    "                lambda record: _decode_record(record, name_to_features),\n",
    "                batch_size=batch_size,\n",
    "                drop_remainder=drop_remainder))\n",
    "\n",
    "        return d\n",
    "\n",
    "    return input_fn\n",
    "\n",
    "\n",
    "def _truncate_seq_pair(tokens_a, tokens_b, max_length):\n",
    "    \"\"\"Truncates a sequence pair in place to the maximum length.\"\"\"\n",
    "\n",
    "    # This is a simple heuristic which will always truncate the longer sequence\n",
    "    # one token at a time. This makes more sense than truncating an equal percent\n",
    "    # of tokens from each, since if one sequence is very short then each token\n",
    "    # that's truncated likely contains more information than a longer sequence.\n",
    "    while True:\n",
    "        total_length = len(tokens_a) + len(tokens_b)\n",
    "        if total_length <= max_length:\n",
    "            break\n",
    "        if len(tokens_a) > len(tokens_b):\n",
    "            tokens_a.pop()\n",
    "        else:\n",
    "            tokens_b.pop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute # train and warmup steps from batch size\n",
    "num_train_steps = int(len(train_examples) / BATCH_SIZE * NUM_TRAIN_EPOCHS)\n",
    "num_warmup_steps = int(num_train_steps * WARMUP_PROPORTION)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_file = os.path.join('./working', \"train.tf_record\")\n",
    "#filename = Path(train_file)\n",
    "if not os.path.exists(train_file):\n",
    "    open(train_file, 'w').close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:***** Running training *****\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0511 14:08:54.477026  4008 tf_logging.py:115] ***** Running training *****\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:  Num examples = 10000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0511 14:08:54.480028  4008 tf_logging.py:115]   Num examples = 10000\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:  Batch size = 32\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0511 14:08:54.481029  4008 tf_logging.py:115]   Batch size = 32\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:  Num steps = 312\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0511 14:08:54.483028  4008 tf_logging.py:115]   Num steps = 312\n"
     ]
    }
   ],
   "source": [
    "file_based_convert_examples_to_features(\n",
    "            train_examples, MAX_SEQ_LENGTH, tokenizer, train_file)\n",
    "tf.logging.info(\"***** Running training *****\")\n",
    "tf.logging.info(\"  Num examples = %d\", len(train_examples))\n",
    "tf.logging.info(\"  Batch size = %d\", BATCH_SIZE)\n",
    "tf.logging.info(\"  Num steps = %d\", num_train_steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_input_fn = file_based_input_fn_builder(\n",
    "    input_file=train_file,\n",
    "    seq_length=MAX_SEQ_LENGTH,\n",
    "    is_training=True,\n",
    "    drop_remainder=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_model(bert_config, is_training, input_ids, input_mask, segment_ids,\n",
    "                 labels, num_labels, use_one_hot_embeddings):\n",
    "    \"\"\"Creates a classification model.\"\"\"\n",
    "    model = modeling.BertModel(\n",
    "        config=bert_config,\n",
    "        is_training=is_training,\n",
    "        input_ids=input_ids,\n",
    "        input_mask=input_mask,\n",
    "        token_type_ids=segment_ids,\n",
    "        use_one_hot_embeddings=use_one_hot_embeddings)\n",
    "\n",
    "    # In the demo, we are doing a simple classification task on the entire\n",
    "    # segment.\n",
    "    #\n",
    "    # If you want to use the token-level output, use model.get_sequence_output()\n",
    "    # instead.\n",
    "    output_layer = model.get_pooled_output()\n",
    "\n",
    "    hidden_size = output_layer.shape[-1].value\n",
    "\n",
    "    output_weights = tf.get_variable(\n",
    "        \"output_weights\", [num_labels, hidden_size],\n",
    "        initializer=tf.truncated_normal_initializer(stddev=0.02))\n",
    "\n",
    "    output_bias = tf.get_variable(\n",
    "        \"output_bias\", [num_labels], initializer=tf.zeros_initializer())\n",
    "\n",
    "    with tf.variable_scope(\"loss\"):\n",
    "        if is_training:\n",
    "            # I.e., 0.1 dropout\n",
    "            output_layer = tf.nn.dropout(output_layer, keep_prob=0.9)\n",
    "\n",
    "        logits = tf.matmul(output_layer, output_weights, transpose_b=True)\n",
    "        logits = tf.nn.bias_add(logits, output_bias)\n",
    "        \n",
    "        # probabilities = tf.nn.softmax(logits, axis=-1) ### multiclass case\n",
    "        probabilities = tf.nn.sigmoid(logits)#### multi-label case\n",
    "        \n",
    "        labels = tf.cast(labels, tf.float32)\n",
    "        tf.logging.info(\"num_labels:{};logits:{};labels:{}\".format(num_labels, logits, labels))\n",
    "        per_example_loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=labels, logits=logits)\n",
    "        loss = tf.reduce_mean(per_example_loss)\n",
    "\n",
    "        # probabilities = tf.nn.softmax(logits, axis=-1)\n",
    "        # log_probs = tf.nn.log_softmax(logits, axis=-1)\n",
    "        #\n",
    "        # one_hot_labels = tf.one_hot(labels, depth=num_labels, dtype=tf.float32)\n",
    "        #\n",
    "        # per_example_loss = -tf.reduce_sum(one_hot_labels * log_probs, axis=-1)\n",
    "        # loss = tf.reduce_mean(per_example_loss)\n",
    "\n",
    "        return (loss, per_example_loss, logits, probabilities)\n",
    "\n",
    "\n",
    "def model_fn_builder(bert_config, num_labels, init_checkpoint, learning_rate,\n",
    "                     num_train_steps, num_warmup_steps, use_tpu,\n",
    "                     use_one_hot_embeddings):\n",
    "    \"\"\"Returns `model_fn` closure for TPUEstimator.\"\"\"\n",
    "\n",
    "    def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument\n",
    "        \"\"\"The `model_fn` for TPUEstimator.\"\"\"\n",
    "\n",
    "        #tf.logging.info(\"*** Features ***\")\n",
    "        #for name in sorted(features.keys()):\n",
    "        #    tf.logging.info(\"  name = %s, shape = %s\" % (name, features[name].shape))\n",
    "\n",
    "        input_ids = features[\"input_ids\"]\n",
    "        input_mask = features[\"input_mask\"]\n",
    "        segment_ids = features[\"segment_ids\"]\n",
    "        label_ids = features[\"label_ids\"]\n",
    "        is_real_example = None\n",
    "        if \"is_real_example\" in features:\n",
    "             is_real_example = tf.cast(features[\"is_real_example\"], dtype=tf.float32)\n",
    "        else:\n",
    "             is_real_example = tf.ones(tf.shape(label_ids), dtype=tf.float32)\n",
    "\n",
    "        is_training = (mode == tf.estimator.ModeKeys.TRAIN)\n",
    "\n",
    "        (total_loss, per_example_loss, logits, probabilities) = create_model(\n",
    "            bert_config, is_training, input_ids, input_mask, segment_ids, label_ids,\n",
    "            num_labels, use_one_hot_embeddings)\n",
    "\n",
    "        tvars = tf.trainable_variables()\n",
    "        initialized_variable_names = {}\n",
    "        scaffold_fn = None\n",
    "        if init_checkpoint:\n",
    "            (assignment_map, initialized_variable_names\n",
    "             ) = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint)\n",
    "            if use_tpu:\n",
    "\n",
    "                def tpu_scaffold():\n",
    "                    tf.train.init_from_checkpoint(init_checkpoint, assignment_map)\n",
    "                    return tf.train.Scaffold()\n",
    "\n",
    "                scaffold_fn = tpu_scaffold\n",
    "            else:\n",
    "                tf.train.init_from_checkpoint(init_checkpoint, assignment_map)\n",
    "\n",
    "        tf.logging.info(\"**** Trainable Variables ****\")\n",
    "        for var in tvars:\n",
    "            init_string = \"\"\n",
    "            if var.name in initialized_variable_names:\n",
    "                init_string = \", *INIT_FROM_CKPT*\"\n",
    "            #tf.logging.info(\"  name = %s, shape = %s%s\", var.name, var.shape,init_string)\n",
    "\n",
    "        output_spec = None\n",
    "        if mode == tf.estimator.ModeKeys.TRAIN:\n",
    "\n",
    "            train_op = optimization.create_optimizer(\n",
    "                total_loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu)\n",
    "\n",
    "            output_spec = tf.estimator.EstimatorSpec(\n",
    "                mode=mode,\n",
    "                loss=total_loss,\n",
    "                train_op=train_op,\n",
    "                scaffold=scaffold_fn)\n",
    "        elif mode == tf.estimator.ModeKeys.EVAL:\n",
    "\n",
    "            def metric_fn(per_example_loss, label_ids, probabilities, is_real_example):\n",
    "\n",
    "                logits_split = tf.split(probabilities, num_labels, axis=-1)\n",
    "                label_ids_split = tf.split(label_ids, num_labels, axis=-1)\n",
    "                # metrics change to auc of every class\n",
    "                eval_dict = {}\n",
    "                for j, logits in enumerate(logits_split):\n",
    "                    label_id_ = tf.cast(label_ids_split[j], dtype=tf.int32)\n",
    "                    current_auc, update_op_auc = tf.metrics.auc(label_id_, logits)\n",
    "                    eval_dict[str(j)] = (current_auc, update_op_auc)\n",
    "                eval_dict['eval_loss'] = tf.metrics.mean(values=per_example_loss)\n",
    "                return eval_dict\n",
    "\n",
    "                ## original eval metrics\n",
    "                # predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)\n",
    "                # accuracy = tf.metrics.accuracy(\n",
    "                #     labels=label_ids, predictions=predictions, weights=is_real_example)\n",
    "                # loss = tf.metrics.mean(values=per_example_loss, weights=is_real_example)\n",
    "                # return {\n",
    "                #     \"eval_accuracy\": accuracy,\n",
    "                #     \"eval_loss\": loss,\n",
    "                # }\n",
    "\n",
    "            eval_metrics = metric_fn(per_example_loss, label_ids, probabilities, is_real_example)\n",
    "            output_spec = tf.estimator.EstimatorSpec(\n",
    "                mode=mode,\n",
    "                loss=total_loss,\n",
    "                eval_metric_ops=eval_metrics,\n",
    "                scaffold=scaffold_fn)\n",
    "        else:\n",
    "            print(\"mode:\", mode,\"probabilities:\", probabilities)\n",
    "            output_spec = tf.estimator.EstimatorSpec(\n",
    "                mode=mode,\n",
    "                predictions={\"probabilities\": probabilities},\n",
    "                scaffold=scaffold_fn)\n",
    "        return output_spec\n",
    "\n",
    "    return model_fn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "OUTPUT_DIR = \"./working/output\"\n",
    "# Specify outpit directory and number of checkpoint steps to save\n",
    "run_config = tf.estimator.RunConfig(\n",
    "    model_dir=OUTPUT_DIR,\n",
    "    save_summary_steps=SAVE_SUMMARY_STEPS,\n",
    "    keep_checkpoint_max=1,\n",
    "    save_checkpoints_steps=SAVE_CHECKPOINTS_STEPS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using config: {'_model_dir': './working/output', '_tf_random_seed': None, '_save_summary_steps': 500, '_save_checkpoints_steps': 1000, '_save_checkpoints_secs': None, '_session_config': None, '_keep_checkpoint_max': 1, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': None, '_device_fn': None, '_service': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x0000026C1A74AEF0>, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0511 14:08:54.546028  4008 tf_logging.py:115] Using config: {'_model_dir': './working/output', '_tf_random_seed': None, '_save_summary_steps': 500, '_save_checkpoints_steps': 1000, '_save_checkpoints_secs': None, '_session_config': None, '_keep_checkpoint_max': 1, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': None, '_device_fn': None, '_service': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x0000026C1A74AEF0>, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1}\n"
     ]
    }
   ],
   "source": [
    "bert_config = modeling.BertConfig.from_json_file(BERT_CONFIG)\n",
    "model_fn = model_fn_builder(\n",
    "  bert_config=bert_config,\n",
    "  num_labels= len(LABEL_COLUMNS),\n",
    "  init_checkpoint=BERT_INIT_CHKPNT,\n",
    "  learning_rate=LEARNING_RATE,\n",
    "  num_train_steps=num_train_steps,\n",
    "  num_warmup_steps=num_warmup_steps,\n",
    "  use_tpu=False,\n",
    "  use_one_hot_embeddings=False)\n",
    "\n",
    "estimator = tf.estimator.Estimator(\n",
    "  model_fn=model_fn,\n",
    "  config=run_config,\n",
    "  params={\"batch_size\": BATCH_SIZE})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Beginning Training!\n",
      "INFO:tensorflow:Calling model_fn.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0511 14:09:03.337898  4008 tf_logging.py:115] Calling model_fn.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:num_labels:6;logits:Tensor(\"loss/BiasAdd:0\", shape=(32, 6), dtype=float32);labels:Tensor(\"loss/Cast:0\", shape=(32, 6), dtype=float32)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0511 14:09:06.024065  4008 tf_logging.py:115] num_labels:6;logits:Tensor(\"loss/BiasAdd:0\", shape=(32, 6), dtype=float32);labels:Tensor(\"loss/Cast:0\", shape=(32, 6), dtype=float32)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:**** Trainable Variables ****\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0511 14:09:06.839105  4008 tf_logging.py:115] **** Trainable Variables ****\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Done calling model_fn.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0511 14:09:16.661103  4008 tf_logging.py:115] Done calling model_fn.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Create CheckpointSaverHook.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0511 14:09:16.666073  4008 tf_logging.py:115] Create CheckpointSaverHook.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Graph was finalized.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0511 14:09:35.468607  4008 tf_logging.py:115] Graph was finalized.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Restoring parameters from ./working/output\\model.ckpt-0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0511 14:09:35.488561  4008 tf_logging.py:115] Restoring parameters from ./working/output\\model.ckpt-0\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Running local_init_op.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0511 14:09:39.940517  4008 tf_logging.py:115] Running local_init_op.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Done running local_init_op.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0511 14:09:40.076634  4008 tf_logging.py:115] Done running local_init_op.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Saving checkpoints for 0 into ./working/output\\model.ckpt.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0511 14:10:10.590216  4008 tf_logging.py:115] Saving checkpoints for 0 into ./working/output\\model.ckpt.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:loss = 0.6658432, step = 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0511 14:18:10.168156  4008 tf_logging.py:115] loss = 0.6658432, step = 1\n"
     ]
    }
   ],
   "source": [
    "print(f'Beginning Training!')\n",
    "current_time = datetime.now()\n",
    "estimator.train(input_fn=train_input_fn, max_steps=num_train_steps)\n",
    "print(\"Training took time \", datetime.now() - current_time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_file = os.path.join('./working', \"eval.tf_record\")\n",
    "#filename = Path(train_file)\n",
    "if not os.path.exists(eval_file):\n",
    "    open(eval_file, 'w').close()\n",
    "\n",
    "eval_examples = create_examples(x_val)\n",
    "file_based_convert_examples_to_features(\n",
    "    eval_examples, MAX_SEQ_LENGTH, tokenizer, eval_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This tells the estimator to run through the entire set.\n",
    "eval_steps = None\n",
    "\n",
    "eval_drop_remainder = False\n",
    "eval_input_fn = file_based_input_fn_builder(\n",
    "    input_file=eval_file,\n",
    "    seq_length=MAX_SEQ_LENGTH,\n",
    "    is_training=False,\n",
    "    drop_remainder=False)\n",
    "\n",
    "result = estimator.evaluate(input_fn=eval_input_fn, steps=eval_steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_eval_file = os.path.join(\"./working\", \"eval_results.txt\")\n",
    "with tf.gfile.GFile(output_eval_file, \"w\") as writer:\n",
    "    tf.logging.info(\"***** Eval results *****\")\n",
    "    for key in sorted(result.keys()):\n",
    "        tf.logging.info(\"  %s = %s\", key, str(result[key]))\n",
    "        writer.write(\"%s = %s\\n\" % (key, str(result[key])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_test = test[:10000] #testing a small sample\n",
    "x_test = x_test.reset_index(drop=True)\n",
    "predict_examples = create_examples(x_test,False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_features = convert_examples_to_features(predict_examples, MAX_SEQ_LENGTH, tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def input_fn_builder(features, seq_length, is_training, drop_remainder):\n",
    "  \"\"\"Creates an `input_fn` closure to be passed to TPUEstimator.\"\"\"\n",
    "\n",
    "  all_input_ids = []\n",
    "  all_input_mask = []\n",
    "  all_segment_ids = []\n",
    "  all_label_ids = []\n",
    "\n",
    "  for feature in features:\n",
    "    all_input_ids.append(feature.input_ids)\n",
    "    all_input_mask.append(feature.input_mask)\n",
    "    all_segment_ids.append(feature.segment_ids)\n",
    "    all_label_ids.append(feature.label_ids)\n",
    "\n",
    "  def input_fn(params):\n",
    "    \"\"\"The actual input function.\"\"\"\n",
    "    batch_size = params[\"batch_size\"]\n",
    "\n",
    "    num_examples = len(features)\n",
    "\n",
    "    # This is for demo purposes and does NOT scale to large data sets. We do\n",
    "    # not use Dataset.from_generator() because that uses tf.py_func which is\n",
    "    # not TPU compatible. The right way to load data is with TFRecordReader.\n",
    "    d = tf.data.Dataset.from_tensor_slices({\n",
    "        \"input_ids\":\n",
    "            tf.constant(\n",
    "                all_input_ids, shape=[num_examples, seq_length],\n",
    "                dtype=tf.int32),\n",
    "        \"input_mask\":\n",
    "            tf.constant(\n",
    "                all_input_mask,\n",
    "                shape=[num_examples, seq_length],\n",
    "                dtype=tf.int32),\n",
    "        \"segment_ids\":\n",
    "            tf.constant(\n",
    "                all_segment_ids,\n",
    "                shape=[num_examples, seq_length],\n",
    "                dtype=tf.int32),\n",
    "        \"label_ids\":\n",
    "            tf.constant(all_label_ids, shape=[num_examples, len(LABEL_COLUMNS)], dtype=tf.int32),\n",
    "    })\n",
    "\n",
    "    if is_training:\n",
    "      d = d.repeat()\n",
    "      d = d.shuffle(buffer_size=100)\n",
    "\n",
    "    d = d.batch(batch_size=batch_size, drop_remainder=drop_remainder)\n",
    "    return d\n",
    "\n",
    "  return input_fn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Beginning Predictions!')\n",
    "current_time = datetime.now()\n",
    "\n",
    "predict_input_fn = input_fn_builder(features=test_features, seq_length=MAX_SEQ_LENGTH, is_training=False, drop_remainder=False)\n",
    "predictions = estimator.predict(predict_input_fn)\n",
    "print(\"Prediction took time \", datetime.now() - current_time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_output(predictions):\n",
    "    probabilities = []\n",
    "    for (i, prediction) in enumerate(predictions):\n",
    "        preds = prediction[\"probabilities\"]\n",
    "        probabilities.append(preds)\n",
    "    dff = pd.DataFrame(probabilities)\n",
    "    dff.columns = LABEL_COLUMNS\n",
    "    \n",
    "    return dff\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_df = create_output(predictions)\n",
    "merged_df =  pd.concat([x_test, output_df], axis=1)\n",
    "submission = merged_df.drop(['comment_text'], axis=1)\n",
    "submission.to_csv(\"sample_submission0.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "submission.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Thanks for Reading, most of the code is from below examples\n",
    "\n",
    "https://github.com/google-research/bert/blob/master/run_classifier.py\n",
    "\n",
    "https://github.com/yajian/bert/blob/master/run_multilabels_classifier.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:tensorflow-bert] *",
   "language": "python",
   "name": "conda-env-tensorflow-bert-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
